package heap_trianning

import (
	"traning/algorithm/utility/heap02"
)

func GetTopKCustomersWithGreaterHeap(customers []int, operate []bool, K int) [][]int {
	ans := make([][]int, 0)
	if K == 0 {
		return ans
	}
	customersMap := make(map[int]*Customer)
	daddys := heap02.NewHeapGreater(K, func(c1, c2 interface{}) bool { // 初始化得奖区的加强堆
		customer1, customer2 := c1.(*Customer), c2.(*Customer)
		if customer1.Buy != customer2.Buy {
			return customer1.Buy < customer2.Buy
		}
		return customer1.Time < customer2.Time
	})
	cands := heap02.NewHeapGreater(len(customers), func(c1, c2 interface{}) bool { // 初始化得奖区的加强堆
		customer1, customer2 := c1.(*Customer), c2.(*Customer)
		if customer1.Buy != customer2.Buy {
			return customer1.Buy > customer2.Buy
		}
		return customer1.Time < customer2.Time
	})
	for time, customerId := range customers {
		if _, ok := customersMap[customerId]; !ok && !operate[time] {
			// 不存在的顾客且发生退货，则直接获取上一次的获奖人，并继续
			ans = append(ans, ToCustomerIds(daddys.GetAllElements()))
			continue
		}
		var curCustomer *Customer
		if c, ok := customersMap[customerId]; !ok {
			curCustomer = &Customer{
				ID: customerId,
			}
			customersMap[customerId] = curCustomer
		} else {
			curCustomer = c
		}
		if operate[time] {
			curCustomer.Buy++
		} else {
			curCustomer.Buy--
		}
		// 清空 customersMap 如果现在顾客的购买量为 0 需要清除
		if curCustomer.Buy == 0 {
			delete(customersMap, curCustomer.ID)
		}
		// 将新客户加入候将区或得奖区
		if !daddys.Contain(curCustomer) && !cands.Contain(curCustomer) {
			if daddys.Size() < K {
				curCustomer.Time = time
				daddys.Add(curCustomer)
			} else {
				curCustomer.Time = time
				cands.Add(curCustomer)
			}
		}
		// 此时顾客一定在两个区中且购买量为0，则需要清除
		if daddys.Contain(curCustomer) && curCustomer.Buy == 0 {
			daddys.Remove(curCustomer)
		}
		if cands.Contain(curCustomer) && curCustomer.Buy == 0 {
			cands.Remove(curCustomer)
		}
		// 元素发生变更需要维护堆
		if daddys.Contain(curCustomer) {
			daddys.Resign(curCustomer)
		}
		if cands.Contain(curCustomer) {
			cands.Resign(curCustomer)
		}
		moveGreaterHeap(&daddys, &cands, time, K)
		ans = append(ans, ToCustomerIds(daddys.GetAllElements()))
	}
	return ans
}

func moveGreaterHeap(daddys, cands *heap02.HeapGreater, time, K int) {
	if cands.Size() == 0 {
		return
	}
	if daddys.Size() < K { // 得奖区有客户退出
		newDaddy := cands.Peek().(*Customer)
		cands.Remove(newDaddy)
		newDaddy.Time = time
		daddys.Add(newDaddy)
	} else {
		if cands.Peek().(*Customer).Buy > daddys.Peek().(*Customer).Buy {
			newDaddy := cands.Peek().(*Customer)
			oldDaddy := daddys.Peek().(*Customer)
			cands.Remove(newDaddy)
			daddys.Remove(oldDaddy)
			newDaddy.Time = time
			oldDaddy.Time = time
			cands.Add(oldDaddy)
			daddys.Add(newDaddy)
		}
	}
}

func ToCustomerIds(daddys []interface{}) []int {
	var daddyIds []int
	for _, c := range daddys {
		daddyIds = append(daddyIds, c.(*Customer).ID)
	}
	return daddyIds
}
